#!/usr/bin/env python
# -*- coding: utf-8 -*-

import socket
import threading
import os
import time
import re
import IPy
import argparse

class Portscan:
    '''
    端口扫描类
    '''

    def __init__(self):
        '''
        创建扫描结果保存元组
        '''
        self.results = {}

    def check_tcp(self, ip, port):
        '''
        tcp扫描
        :param ip:
        :param port:
        :return:
        '''
        sock_tcp = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
        # 设置超时
        sock_tcp.settimeout(10)
        status = sock_tcp.connect_ex((ip, port))
        # print(status)
        if status == 0:
            self.results[ip][port] = '[+]{0}:{1}  open '.format(ip, port)
            # print('[+]{0}:{1} open '.format(ip, port))
            sock_tcp.close()

    def deal_range(self, range):
        '''
        正则匹配用户输入的port_range,允许范围0~65535
        :param range
        :return: 端口范围begin和end
        '''
        try:
            begin = int(range)
            end = int(range) + 1
            if (int(begin) < 0 or int(end) > 65535):
                print('port number wrong...')
                os._exit(0)
        except:
            begin = re.search(r'(.*)-(.*)', range).group(1)
            end = int(re.search(r'(.*)-(.*)', range).group(2))+1
            # print(begin, end)
            if int(begin) < 0 or int(end) > 65535:
                print('port number wrong...')
                os._exit(0)
        return begin, end

    def is_ip(self,address):
        '''
        对ip进行判断
        :param:ip
        :return: True or False
        '''
        try:
            IPy.IP(address)
            return True

        except Exception as e:
            return False

    def deal_ip(self, ip):
        '''
        正则匹配用户输入的ip
        :param ip
        :return: 用户指定的ip范围列表
        '''
        ips = []
        if re.search(r'(.*)-(.*)', ip):
            ip_split = re.search(r'(.*)-(.*)', ip).group(1).split('.')
            ip_end = re.search(r'(.*)-(.*)', ip).group(2)
            for x in range(int(ip_split[3]), int(ip_end) + 1):
                ip_split[3] = str(x)
                ip = ('.').join(ip_split)
                # print(ip)
                if portscan.is_ip(ip):
                    ips.append(ip)

        else:
            if portscan.is_ip(ip):
                ips.append(ip)
        self.ips = ips
        return ips

    def run(self, ip, begin, end):
        self.results[ip] = {}
        '''
        对指定ip进行端口进行扫描：开启tcp多端口同时扫描
        :param ip:
        :param begin:
        :param end:
        '''

        print(ip, 'start...')
        print("[!]working hard...")

        tcp_threads = []
        for port in range(begin, end):
            t = threading.Thread(target=self.check_tcp, args=(ip, port))
            tcp_threads.append(t)

        # for thread in tcp_threads:
        #     thread.start()
        # for thread in tcp_threads:
        #     thread.join()
  
        for thread in tcp_threads:
            thread.setDaemon(True)
            thread.start()

        for thread in tcp_threads:
            thread.join()


if __name__ == '__main__':

    # 获取命令行参数
    parser = argparse.ArgumentParser()
    parser.add_argument("-H", "--target", help="Your target ip, eg:10.10.10.130 or 10.10.10.130-144", required=True)
    parser.add_argument("-P", "--port_range", help="Port range in 0~65535, eg:80 or 80-500", required=True)
    args = parser.parse_args()

    # 计时
    time_start = time.time()

    # 创建端口扫描类实例
    portscan = Portscan()

    # 获取端口号
    begin, end = portscan.deal_range(args.port_range)

    # 获取ip
    ips = portscan.deal_ip(args.target)

    # 输入ip错误报错处理
    if not ips:
        print('All ip Unreachable or packet loss!')
        os._exit(0)

    # 创建进程池
    # pool = ThreadPoolExecutor(255)

    # 创建线程列表，同时对多个ip进行扫描
    threads=[]
    # 对每个IP创建一个线程
    for ip in ips:
        t = threading.Thread(target=portscan.run, args=(ip, int(begin), int(end)))
        threads.append(t)

    for thread in threads:
        thread.setDaemon(True)
        thread.start()

    for thread in threads:
        thread.join()

    # 将TCP扫描结果写入portscan.results中
    #     for port in range(int(begin), int(end)):
    #         portscan.results[ip][port] = '[+]{0}:{1}/tcp open '.format(ip, port)

    # 将portscan.results中的结果写入文件，并生成
    for ip in portscan.results:
        with open(ip + '.txt', 'w') as f:
            for key in portscan.results[ip]:
                print(portscan.results[ip][key])
                f.write(portscan.results[ip][key] + '\n')

    # 计时结束
    time_end = time.time()

    # 运行时间输出
    print('time cost:', "{:.5}".format(time_end - time_start), 'done.')